package cn.iocoder.yudao.module.engine.service.core.messagestore.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.iocoder.yudao.module.engine.model.context.AIMessageContext;
import cn.iocoder.yudao.module.engine.model.context.ControlEngineContext;
import cn.iocoder.yudao.module.engine.model.context.EngineSessionContext;
import cn.iocoder.yudao.module.engine.model.engine.EngineOutputData;
import cn.iocoder.yudao.module.engine.service.core.messagestore.AbstractMessageStore;
import cn.iocoder.yudao.module.engine.service.core.messagestore.MessageStoreManager;
import cn.iocoder.yudao.module.resource.api.IMessageCollectionApi;
import cn.iocoder.yudao.module.resource.api.model.message.AudioDataDTO;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Service;

import javax.annotation.Nullable;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import java.util.List;
import java.util.Map;

@Slf4j
@Order(103)
@Service("messageStore-vector")
public class VectorMessageStore extends AbstractMessageStore {
    @Resource
    private MessageStoreManager messageStoreManager;
    @Resource
    private IMessageCollectionApi messageVectorApi;

    @PostConstruct
    public void init() {
        messageStoreManager.registerStore("messageStore-vector", this);
    }

    @Override
    public void executeSaveMessages(ControlEngineContext engineContext,
                                    @Nullable AIMessageContext messageCtx,
                                    EngineOutputData.Audio audio) {
        if (messageCtx != null && messageCtx.isFromTTSCache()) {
            log.info("消息来自TTS缓存, 不保存到向量库");
            return;
        }

        saveMsgToVector(engineContext.getSessionCtx(), messageCtx, audio);
    }

    private void saveMsgToVector(EngineSessionContext sessionContext,
                                 @Nullable AIMessageContext messageCtx,
                                 EngineOutputData.Audio audio) {
        try {
            final String messageId = audio.getMessageId();
            final String userText = audio.getFromMessage();
            final String replyText = getReplyText(messageCtx, audio);

            final List<AudioDataDTO> audioDataList = getAudioDataList(messageCtx, audio);
            if (CollUtil.isEmpty(audioDataList)) {
                log.warn("消息没有音频, 不保存到向量库, messageId = {}, userText = {}", messageId, userText);
                return;
            }

            if (messageCtx != null && messageCtx.isFromLLMCache()) {
                log.info("消息来自LLM缓存, 更新向量库, messageId = {}, userText = {}", messageId, userText);
                final Map<String, List<? extends Number>> vectorValues = messageVectorApi.getVectorValues(userText, replyText);
                messageVectorApi.updateMessage(
                        sessionContext.getToneId(), audio.getMessageSendType(), messageId,
                        audioDataList, userText, replyText, vectorValues
                );
            } else {
//                final MessageCollectionResDTO messageVectorRes = messageVectorApi.queryMessage(0.9f, 1, userText, replyText);
//                if (CollUtil.isNotEmpty(messageVectorRes.getMessages())) {
//                    log.info("消息已存在, 不保存到向量库, messageId = {}, userText = {}", messageId, userText);
//                    return;
//                }
//
//                log.info("消息不存在, 保存到向量库, messageId = {}, userText = {}", messageId, userText);
//                messageVectorApi.insertMessage(
//                        sessionContext.getToneId(), audio.getMessageSendType(), messageId,
//                        audioDataList, userText, replyText, messageVectorRes.getVectorValues()
//                );
            }
        } catch (Exception e) {
            log.error(StrUtil.format("保存消息到向量库失败, messageId = {}, userText = {}",
                    audio.getMessageId(), audio.getFromMessage()), e);
        }
    }


}
